// Copyright (c) 2025 Yao Yu
// This Source Code Form is subject to the terms of the Mozilla Public
// License, v. 2.0. If a copy of the MPL was not distributed with this
// file, You can obtain one at https://mozilla.org/MPL/2.0/.

use crossbeam::queue::SegQueue;
use alloc::{sync::{Weak, Arc}, vec};
use bitflags::bitflags;
use core::sync::atomic::{AtomicU8, AtomicUsize, Ordering};
use crate::{
    context::Context,
    prelude::*,
    thread::{EndPoint, IPC_BUF_SIZE, Thread},
    utils::{min_usize, random::randu32}
};
use hashbrown::HashMap;
use hmac::{Hmac, Mac};
use lazy_static::*;
use sha2::Sha256;
use ostd::{
    cpu::context::UserContext,
    mm::{FallibleVmRead, FallibleVmWrite, VmSpace, VmWriter, VmReader},
    sync::{Mutex, RwLock, WaitQueue},
};

// 全局端点哈希表， lazy_static在第一次使用时加载。
lazy_static! {
    static ref GLOBAL_ENDPOINTS: RwLock<HashMap<u64, Arc<EndPoint>>> = RwLock::new(HashMap::new());
    static ref ENDPOINT_ID: AtomicUsize = AtomicUsize::new(0);
    // static ref WAIT_QUEUE: WaitQueue = WaitQueue::new();
}

bitflags!{
    #[derive(Clone, Copy)]
    pub struct Capability: u8{
        const CAP_NULL     = 0b00000000;
        const CAP_SEND     = 0b00000001;
        const CAP_RECEIVE  = 0b00000010;
        const CAP_GRANT    = 0b00000100;
        const CAP_DELETE   = 0b00001000;
        const CAP_NOTIFY   = 0b00010000;
        const CAP_READ     = 0b00100000;
        const CAP_WRITE    = 0b01000000;
        const CAP_ALL = Self::CAP_SEND.bits()
                    | Self::CAP_RECEIVE.bits()
                    | Self::CAP_GRANT.bits()
                    | Self::CAP_DELETE.bits()
                    | Self::CAP_NOTIFY.bits()
                    | Self::CAP_READ.bits()
                    | Self::CAP_WRITE.bits();
    }
}

#[derive(PartialEq)]
#[repr(u8)]
enum IpcStatus {
    Init = 0,
    Ready = 1,
    Blocked = 2,
}
impl Into<u8> for IpcStatus {
    fn into(self) -> u8 {
        self as u8
    }
}

// TODO:如果能通过asterinas中宏的方式derive TryFromInt的话就更好了。
impl From<u8> for IpcStatus {
    fn from(value: u8) -> Self {
        match value {
            0 => IpcStatus::Init,
            1 => IpcStatus::Ready,
            2 => IpcStatus::Blocked,
            _ => panic!("Invalid IpcStatus value"),
        }
    }
}

// ipc通信需要分配端点
pub struct EndPoint {
    // 所属线程
    owner: Weak<Thread>,
    // 记录发送线程的队列
    send_queue: Mutex<SegQueue<Weak<Thread>>>,
    // 能力令牌
    capability: Capability,

    // ipc阻塞状态，用于恢复被阻塞线程
    sender_ipc_status: AtomicU8,
    receiver_ipc_status: AtomicU8,

    // 存放因接收方未就绪而阻塞的发送线程
    send_wait_queue: WaitQueue,
    // 存放因发送方未就绪而阻塞的接收线程
    receive_wait_queue: WaitQueue,
}

impl EndPoint {
    fn new(thread: Weak<Thread>, capability: Capability) -> Self {
        Self {
            owner: thread,
            send_queue: Mutex::new(SegQueue::new()),
            capability,
            sender_ipc_status: AtomicU8::new(IpcStatus::Init.into()),
            receiver_ipc_status: AtomicU8::new(IpcStatus::Init.into()),
            send_wait_queue: WaitQueue::new(),
            receive_wait_queue: WaitQueue::new(),
        }
    }

    // send时，将发送线程放入队列
    fn push_thread2queue(&self, thread: Weak<Thread>) {
        self.send_queue.lock().push(thread);
    }

    // receive时，从队列中取出发送线程
    fn pop_thread_from_queue(&self) -> Option<Weak<Thread>> {
        self.send_queue.lock().pop()
    }

    // 获得令牌
    fn get_capability_as_u8(&self) -> u8 {
        self.capability.bits()
    }

    fn contains_cap_read(&self) -> bool {
        self.capability.contains(Capability::CAP_READ)
    }

    // IPC阻塞线程状态相关

    fn set_sender_blocked(&self) {
        self.sender_ipc_status.store(IpcStatus::Blocked.into(), Ordering::Release);
    }

    fn set_sender_ready(&self) {
        self.sender_ipc_status.store(IpcStatus::Ready.into(), Ordering::Release);
    }

    fn sender_ipc_status(&self) -> IpcStatus {
        IpcStatus::from(self.sender_ipc_status.load(Ordering::Acquire))
    }
    fn sender_is_waiting(&self) -> bool {
        self.sender_ipc_status() == IpcStatus::Blocked
    }
    fn sender_is_not_waiting(&self) -> bool {
        self.sender_ipc_status() == IpcStatus::Ready
    }

    fn set_receiver_blocked(&self) {
        self.receiver_ipc_status.store(IpcStatus::Blocked.into(), Ordering::Release);
    }

    fn set_receiver_ready(&self) {
        self.receiver_ipc_status.store(IpcStatus::Ready.into(), Ordering::Release);
    }

    fn receiver_ipc_status(&self) -> IpcStatus {
        IpcStatus::from(self.receiver_ipc_status.load(Ordering::Acquire))
    }
    fn receiver_is_waiting(&self) -> bool {
        self.receiver_ipc_status() == IpcStatus::Blocked
    }
    fn receiver_is_not_waiting(&self) -> bool {
        self.receiver_ipc_status() == IpcStatus::Ready
    }

    // 发送方主动进入睡眠状态，直到接收方唤醒
    fn sender_sleep(&self) {
        self.send_wait_queue.wait_until(|| {
            let status = self.sender_ipc_status();
            if status == IpcStatus::Ready {
                Some(status)
            } else {
                None
            }
        });
    }
    // 接收方主动进入睡眠状态，知道发送方唤醒
    fn receiver_sleep(&self) {
        self.receive_wait_queue.wait_until(|| {
            let status = self.receiver_ipc_status();
            if status == IpcStatus::Ready {
                Some(status)
            } else {
                None
            }
        });
    }
}

// 生成安全句柄(在哈希表中，端点的索引是安全句柄，同时会将安全句柄返回给用户进程)
type HmacSha256 = Hmac<Sha256>;
fn generate_secure_handle(endpoint: Arc<EndPoint>, secret_key: &[u8; 32]) -> u64 {
    // 输入：object（内核对象指针） + secret_key（线程密钥）
    // 输出：64 位加密句柄

    // TODO: 精简数据结构和代码实现
    // 1. 构造待加密数据（对象指针 + 权限 + 随机盐）
    #[repr(C)]
    struct HandleData {
        object: usize,
        rights: u32,
        salt: u32,
    }

    let data = HandleData {
        object: Arc::into_raw(Arc::clone(&endpoint)) as usize,
        rights: endpoint.get_capability_as_u8() as u32,
        salt: randu32(),
    };

    // 2. 使用 HMAC-SHA256 计算哈希
    let input_data = {
        let mut input_data = vec::Vec::new();
        input_data.extend_from_slice(&data.object.to_ne_bytes());
        input_data.extend_from_slice(&data.rights.to_ne_bytes());
        input_data.extend_from_slice(&data.salt.to_ne_bytes());
        input_data
    };

    let mut mac = HmacSha256::new_from_slice(secret_key)
        .expect("HMAC key length invalid");
    mac.update(&input_data);
    let hmac_result = mac.finalize().into_bytes();

    // 3. 取哈希值的前 64 位作为句柄（可根据需求调整长度）
    let handle = u64::from_le_bytes(hmac_result[..8].try_into().unwrap());

    // TODO: 4. 记录句柄到内核的映射表（仅供验证用）
    GLOBAL_ENDPOINTS.write().insert(handle, endpoint);

    handle
}

// 解密句柄，根据句柄从全局哈希表中获得端点
fn decrypt_handle(ep_cap: u64, _secret_key: &[u8; 32]) -> Option<Arc<EndPoint>> {
    // 1. 检查句柄是否在合法范围内（可选）TODO
    

    // 2. 从全局句柄表中查找记录（维护 ep_cap → endpoint* 的映射）
    let endpoint = GLOBAL_ENDPOINTS.read().get(&ep_cap).map(|ep| ep.clone());

    // 3. 验证句柄的 HMAC 是否匹配（防篡改）TODO
    

    // 4. 返回解密后的 endpoint* TODO
    

    endpoint
}

// 创建端点
fn sys_create_endpoint(current_thread: Arc<Thread>) -> isize {
    // 创建端点，赋予发送和接收能力
    let endpoint = Arc::new(
                                            EndPoint::new(
                                                Arc::downgrade(&current_thread), 
                                                Capability::CAP_SEND | Capability::CAP_RECEIVE
                                            )
                                    );
    // 生成安全句柄，并且这里已经将句柄->endpoint插入到全局哈希表中了。
    let ep_cap = generate_secure_handle(endpoint, current_thread.secret_key());
    current_thread.register_cap(ep_cap);

    // 这对吗？
    ep_cap as isize
}

fn sys_close_endpoint() -> isize {
    0
}

const MSG_LEN: usize = 64; // 64 bytes

pub fn sys_send(endpoint: EndPoint, msg: *const u8, ctx: &Context, user: &UserContext) -> Result<isize> {
    let dest_thread = ;
    let caller = ctx.thread();
    // check endpoint capability;
    if !dest_thread.is_recving_from(caller) {
        caller.set_ipc_state(ThredStatus::Running);
        // 加入等待队列
        block();
        return Ok(0);
    }
    copy_message(caller, dest_thread, msg);
    dest_thread.set_recving_false();
    wakeup(dest_thread);
    if caller.send_rec() {
        block();
    }
    Ok(0)
}

// ipc发送系统调用
fn sys_send(ep_cap: u64, msg: *const u8, len: usize, vm_space: &VmSpace, current_thread: Arc<Thread>) -> isize{
    // 1 验证目标端点和能力
    // TODO: 返回-EPERM
    let dest = decrypt_handle(ep_cap, current_thread.secret_key()).unwrap();
    if !dest.contains_cap_read() {
        // 这里锁住调用contains_cap_read()后会立刻释放
        return -1;
    }

    // 检查拷贝长度是否大于最大长度
    if len > IPC_BUF_SIZE {
        return -1;
    }

    // 2 拷贝消息到内核缓冲区
    let buffer = {
        let mut buf = vec![0u8; len];
        let mut reader = vm_space.reader(msg as usize, len).unwrap();
        reader.read_fallible(&mut VmWriter::from(&mut buf as &mut [u8])).unwrap();
        buf
    };

    // 3 检查接收方状态
    let receiver = dest.owner.upgrade().unwrap();
    if dest.receiver_is_waiting() {
        // 如果接收线程正在等待，那就直接将消息放入接收方的缓冲区，顺便唤醒该线程
        receiver.copy2ipc_buf(&buffer);
        // receiver.set_ready();

        // 设置成ready后自动唤醒那个线程了。
        dest.set_receiver_ready();
        // TODO： 如果能够唤醒接收方线程就更好了。现在无法唤醒要靠调度器来将时间片轮转到receiver上。
        // schedule2(receiver);
    } else {
        // 如果接收线程没有等待。那么需要将当前线程flag为等待状态，并将消息放入发送队列
        // notice: 这里的发送方没有将buffer中内容拷贝到接收方的缓冲区，而是将消息放入发送队列。
        current_thread.copy2ipc_buf(&buffer);
        // dest.send_queue.push(Arc::downgrade(&current_thread));
        dest.push_thread2queue(Arc::downgrade(&current_thread));
        // 当前进程释放CPU，切换到其他CPU上来工作
        dest.set_sender_blocked();
        dest.sender_sleep();
    }
    0
}

fn sys_receive(ep_cap: u64, msg: *mut u8, max_len: usize, vm_space: &VmSpace, current_thread: Arc<Thread>) -> isize {
    // 从哈希表中获得端点，遇到None直接返回-1
    let src = match decrypt_handle(ep_cap, current_thread.secret_key()) {
        Some(src) => src,
        None => return -1,
    };
    // let mut srcguard = src.lock();

    let senderop = src.pop_thread_from_queue();
    match senderop {
        // 如果发送队列中有消息，则将消息拷贝到用户缓冲区
        Some(sender) => {
            let sender = match sender.upgrade() {
                Some(sender) => sender,
                None => return -1,
            };

            let message = sender.borrow_ipc_buf();
            let len = min_usize(sender.ipc_msg_size(), max_len);
            let mut writer = vm_space.writer(msg as usize, len).unwrap();
            // 拷贝到用户内存
            writer.write_fallible(&mut VmReader::from(&*message as &[u8])).unwrap();
            // 接收者已经拿到数据了，所以可以开始运行了。
            src.set_sender_ready();
        }
        None => {
            // current_thread->waiting_on = ep;
            // strong TODO : 这里怎么设置waiting_on啊？？
            // current_thread.set_waiting_on(Arc::downgrade(&src));

            src.set_receiver_blocked();
            // TODO: 这里应该将当前线程放入等待队列，等待ep的通知。
            src.receiver_sleep();
            let message = current_thread.borrow_ipc_buf();
            let len = min_usize(current_thread.ipc_msg_size(), max_len);
            let mut writer = vm_space.writer(msg as usize, len).unwrap();
            writer.write_fallible(&mut VmReader::from(&*message as &[u8])).unwrap();
            // 唤醒当前线程，让它开始运行
        }
    }
    0
}

fn sys_notify() -> isize {
    0
}